#pragma once

#include <utility>
#include <functional>
#include <map>
#include <list>
#include <exception>

#include "_socket.h"
#include "conn.h"


class TcpServer {

public:

    explicit TcpServer(int port, int family = IpFamily::v4) : port_(port) {
        WSAData wsData{};
        if (WSAStartup(MAKEWORD(2, 2), &wsData) != 0) {
            throw std::runtime_error("WSAStartup failed: " + std::to_string(WSAGetLastError()));
        }

        serverSocket_ = socket(family, SOCK_STREAM, 0);
        if (serverSocket_ == Err::InvalidSocket) {
            throw std::runtime_error("Could not create socket: " + std::to_string(WSAGetLastError()));
        }

        sockaddr_in serverAddr{};
        serverAddr.sin_family = family;
        serverAddr.sin_addr.s_addr = INADDR_ANY;
        serverAddr.sin_port = htons(port);

        if (bind(serverSocket_, (struct sockaddr *) &serverAddr, sizeof(serverAddr)) == Err::SocketError) {
            throw std::runtime_error("Bind failed with error code: " + std::to_string(WSAGetLastError()));
        }

        if (listen(serverSocket_, SOMAXCONN) == Err::SocketError) {
            throw std::runtime_error("Listen failed with error: " + std::to_string(WSAGetLastError()));
        }

        std::cout << "tcp server listen on " << port_ << std::endl;

        if (!SetNonBlock(serverSocket_, &nonBlock)) {
            throw std::runtime_error("ioctlsocket() failed with error: " + std::to_string(WSAGetLastError()));
        }
    }

    void Run() {
        Select();
    }

    template<class T>
    void Register(int16_t proto, HandlerFunc<T> handler) {
        if (running_) {
            throw std::runtime_error("Cannot register after Run()");
        }
        const auto &iter = handlers_.find(proto);
        if (iter == handlers_.end()) {
            T *instance = new T();
            handlers_[proto] = [instance, handler](TcpConn *conn, const char *data, int size) {
                instance->ParseFromArray(data, size);
                std::string json;
                if(google::protobuf::util::MessageToJsonString(*instance, &json) == google::protobuf::util::Status::OK) {
                    std::cout << "recv " << typeid(*instance).name() << " " << json << std::endl;
                }
                handler(conn, *instance);
            };
        } else {
            throw std::logic_error("duplicated register " + std::to_string(proto));
        }
    }

private:
    void Select() {
        fd_set readSet;
        fd_set writeSet;
        int ret;

        while (1) {
            FD_ZERO(&readSet);
            FD_ZERO(&writeSet);
            FD_SET(serverSocket_, &readSet);

            for (TcpConn *c : conns_) {
                FD_SET(c->Socket(), &writeSet);
                FD_SET(c->Socket(), &readSet);
            }

            // Select
            ret = select(0, &readSet, &writeSet, nullptr, nullptr);
            if (ret == Err::SocketError) {
                throw std::runtime_error("Select Error: " + std::to_string(WSAGetLastError()));
            }

            // Accept
            if (FD_ISSET(serverSocket_, &readSet)) {
                SOCKET sock = accept(serverSocket_, nullptr, nullptr);
                if (sock == Err::InvalidSocket) {
                    throw std::runtime_error("Accept Error: " + std::to_string(WSAGetLastError()));
                }

                if (!SetNonBlock(serverSocket_, &nonBlock)) {
                    throw std::runtime_error("ioctlsocket() failed with error: " + std::to_string(WSAGetLastError()));
                }

                CreateConn(sock);
            }

            // Loop
            auto it = conns_.begin();
            while (it != conns_.end()) {
                TcpConn *conn = *it;
                if (FD_ISSET(conn->Socket(), &readSet)) {
                    static char *BUF = new char[MAX_PACK_SIZE];
                    int16_t proto = 0;
                    int size = 0;
                    if ((ret = conn->Read(&proto, BUF, &size)) == 0) {
                        HandleData(conn, proto, BUF, size);
                    } else {
                        std::cout << GetErrorMsg(ret) << std::endl;
                        it = conns_.erase(it);
                        delete conn;
                        continue;
                    }
                }
                ++it;
            }
        }
    }

    void CreateConn(SOCKET s) {
        _Socket *so = new _Socket();
        so->s = s;
        TcpConn *con = new TcpConn(so);
        conns_.push_back(con);
    }

    void HandleData(TcpConn *con, int proto, const char *data, int size) {
        auto h = handlers_.find(proto);
        if (h == handlers_.end()) {
            std::cout << "Error: " << "Cannot find the handler of proto " << proto << std::endl;
            return;
        }
        con->SetProto(proto);
        h->second(con, data, size);
    }

private:
    std::map<int16_t, std::function<void(TcpConn *, const char *, int)>> handlers_;
    bool running_ = false;
    int port_;
    u_long nonBlock = 1;
    SOCKET serverSocket_;
    std::list<TcpConn *> conns_;
};
